package com.ben.multiple.shard;

import com.ben.multiple.shard.springjdbc.mp.Order;
import com.zaxxer.hikari.HikariDataSource;
import org.apache.shardingsphere.driver.api.ShardingSphereDataSourceFactory;
import org.apache.shardingsphere.infra.config.algorithm.AlgorithmConfiguration;
import org.apache.shardingsphere.infra.config.mode.ModeConfiguration;
import org.apache.shardingsphere.infra.config.props.ConfigurationPropertyKey;
import org.apache.shardingsphere.infra.config.rule.RuleConfiguration;
import org.apache.shardingsphere.mode.repository.standalone.StandalonePersistRepositoryConfiguration;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.strategy.keygen.KeyGenerateStrategyConfiguration;
import org.apache.shardingsphere.sharding.api.config.strategy.sharding.ShardingStrategyConfiguration;
import org.apache.shardingsphere.sharding.api.config.strategy.sharding.StandardShardingStrategyConfiguration;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Random;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

public class RawJdbcTest {
    static DataSource dataSource;

    static {
        Configuration config = new Configuration();
        try {
            dataSource = config.createDataSource();
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }

    public static void main(String[] args) {
        insert();
//        select();
    }


    static AtomicInteger atomicInteger = new AtomicInteger(LocalDateTime.now().getSecond());
    static AtomicLong atomicLong = new AtomicLong(LocalDateTime.now().getMinute());
    static String[] cs = {"A", "B", "C"};

    private static void insert() {
        List<Order> list = new ArrayList<>();
        list.add(new Order(null, new Random().nextInt(2), atomicInteger.getAndIncrement(), atomicLong.getAndIncrement(), cs[new Random().nextInt(3)]));
        list.add(new Order(null, new Random().nextInt(2), atomicInteger.getAndIncrement(), atomicLong.getAndIncrement(), cs[new Random().nextInt(3)]));
        list.add(new Order(null, new Random().nextInt(2), atomicInteger.getAndIncrement(), atomicLong.getAndIncrement(), cs[new Random().nextInt(3)]));
        String sql = "INSERT INTO t_order (order_type, user_id, address_id, status) VALUES(?, ?, ?, ?),(?, ?, ?, ?),(?, ?, ?, ?)";
        try (Connection connection = dataSource.getConnection();
             PreparedStatement preparedStatement = connection.prepareStatement(sql, Statement.RETURN_GENERATED_KEYS)) {
            Order o1 = list.get(0);
            Order o2 = list.get(1);
            Order o3 = list.get(2);
            preparedStatement.setInt(1, o1.getOrderType());
            preparedStatement.setInt(2, o1.getUserId());
            preparedStatement.setLong(3, o1.getAddressId());
            preparedStatement.setString(4, o1.getStatus());

            preparedStatement.setInt(1 + 4, o2.getOrderType());
            preparedStatement.setInt(2 + 4, o2.getUserId());
            preparedStatement.setLong(3 + 4, o2.getAddressId());
            preparedStatement.setString(4 + 4, o2.getStatus());

            preparedStatement.setInt(1 + 8, o3.getOrderType());
            preparedStatement.setInt(2 + 8, o3.getUserId());
            preparedStatement.setLong(3 + 8, o3.getAddressId());
            preparedStatement.setString(4 + 8, o3.getStatus());

            preparedStatement.executeUpdate();
            try (ResultSet resultSet = preparedStatement.getGeneratedKeys()) {
                List<Long> orderIdList = new ArrayList<>();
                while (resultSet.next()) {
                    orderIdList.add(resultSet.getLong(1));
                }
                int size = orderIdList.size();
                for (int i = 0; i < size; i++) {
                    list.get(i).setOrderId(orderIdList.get(i));
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        list.forEach(System.out::println);
        System.out.println("==========================================");
    }

    private static void select() {
        String sql = "SELECT order_id, order_type, user_id, address_id, status FROM t_order";
        List<Order> result = new LinkedList<>();
        try (Connection connection = dataSource.getConnection();
             PreparedStatement preparedStatement = connection.prepareStatement(sql);
             ResultSet resultSet = preparedStatement.executeQuery()) {
            while (resultSet.next()) {
                Order order = new Order();
                order.setOrderId(resultSet.getLong(1));
                order.setOrderType(resultSet.getInt(2));
                order.setUserId(resultSet.getInt(3));
                order.setAddressId(resultSet.getLong(4));
                order.setStatus(resultSet.getString(5));
                result.add(order);
            }
        }catch (Exception e){
            e.printStackTrace();
        }
        result.forEach(System.out::println);
    }

    static final class Configuration {
        private static final int PORT = 3306;
        private static final String USER_NAME = "root";
        private static final String PASSWORD = "q2024E=t!";

        public DataSource createDataSource() throws SQLException {
            return ShardingSphereDataSourceFactory.createDataSource(createModeConfiguration(), createDataSourceMap(), createRuleConfiguration(), createProperties());
        }

        private static ModeConfiguration createModeConfiguration() {
            return new ModeConfiguration("Standalone", new StandalonePersistRepositoryConfiguration("JDBC", new Properties()));
        }

        private Map<String, DataSource> createDataSourceMap() {
            Map<String, DataSource> result = new LinkedHashMap<>();
            result.put("ds_0", createDataSource("192.168.229.146", "raw_jdbc"));
            result.put("ds_1", createDataSource("192.168.229.147", "raw_jdbc"));
            return result;
        }

        private DataSource createDataSource(String ip, String dataSourceName) {
            HikariDataSource result = new HikariDataSource();
            result.setDriverClassName("com.mysql.cj.jdbc.Driver");
            result.setJdbcUrl(String.format("jdbc:mysql://%s:%s/%s?serverTimezone=UTC&useSSL=false&useUnicode=true&characterEncoding=UTF-8&allowPublicKeyRetrieval=true", ip, PORT, dataSourceName));
            result.setUsername(USER_NAME);
            result.setPassword(PASSWORD);
            return result;
        }

        private Collection<RuleConfiguration> createRuleConfiguration() {
            Collection<RuleConfiguration> result = new LinkedList<>();
            result.add(createShardingRuleConfiguration());
            return result;
        }

        private ShardingRuleConfiguration createShardingRuleConfiguration() {
            ShardingRuleConfiguration result = new ShardingRuleConfiguration();
            result.getTables().add(getOrderTableRuleConfiguration());
            result.getShardingAlgorithms().putAll(buildShardingAlgorithms());
            result.getKeyGenerators().putAll(buildKeyGenerators());
            return result;
        }

        private static ShardingTableRuleConfiguration getOrderTableRuleConfiguration() {
            ShardingStrategyConfiguration databaseShardingStrategy = new StandardShardingStrategyConfiguration("user_id", "ds_order_sa");
            KeyGenerateStrategyConfiguration keyGenerateStrategy = new KeyGenerateStrategyConfiguration("order_id", "snowflake_generator");
            ShardingTableRuleConfiguration shardingTableRuleConfiguration = new ShardingTableRuleConfiguration("t_order", "ds_$->{0..1}.t_order");
            shardingTableRuleConfiguration.setDatabaseShardingStrategy(databaseShardingStrategy);
            shardingTableRuleConfiguration.setKeyGenerateStrategy(keyGenerateStrategy);
            return shardingTableRuleConfiguration;
        }

        private static Map<String, AlgorithmConfiguration> buildShardingAlgorithms() {
            Map<String, AlgorithmConfiguration> shardingAlgorithms = new LinkedHashMap<>();
            Properties properties = new Properties();
            properties.put("algorithm-expression", "ds_${user_id % 2}");
            shardingAlgorithms.put("ds_order_sa", new AlgorithmConfiguration("INLINE", properties));
            return shardingAlgorithms;
        }

        private static Map<String, AlgorithmConfiguration> buildKeyGenerators() {
            Map<String, AlgorithmConfiguration> keyGenerators = new LinkedHashMap<>();
            keyGenerators.put("snowflake_generator", new AlgorithmConfiguration("SNOWFLAKE", new Properties()));
            return keyGenerators;
        }

        private Properties createProperties() {
            Properties result = new Properties();
            result.setProperty(ConfigurationPropertyKey.SQL_SHOW.getKey(), "true");
            return result;
        }
    }

}
